require 'spec_helper'

RSpec.describe Msf::Exploit::Retry do

  def create_exploit(info = {})
    mod = Msf::Exploit.allocate
    mod.extend described_class
    mod.send(:initialize, info)
    mod
  end

  describe '#retry_until_truthy' do
    subject do
      create_exploit
    end

    # Quick workaround for Timecop not supporting Process.clock_gettime
    # Timecop feature request: https://github.com/travisjeffery/timecop/issues/220
    let(:mock_clock) do
      clazz = Class.new do
        def initialize
          @current_time = 0
        end

        def now
          @current_time
        end

        def time_travel(new_time)
          @current_time = new_time
        end
      end

      clazz.new
    end

    def increment_time_by(seconds)
        mock_clock.time_travel(mock_clock.now + seconds)
    end

    def expect_sleep_calls(calls)
      expect(subject).to have_received(:sleep).exactly(calls.length).times
      calls.each do |value|
        expect(subject).to have_received(:sleep).with(value)
      end
    end

    before(:each) do
      allow(Process).to receive(:clock_gettime).with(Process::CLOCK_MONOTONIC, :second) { |*_args| mock_clock.now }
      allow(subject).to receive(:sleep) { |seconds| increment_time_by(seconds) }
    end

    context 'when the timeout is negative' do
      it 'does not yield control' do
        result = nil
        expect do |block|
          result = subject.retry_until_truthy(timeout: -1, &block)
        end.to_not yield_control
        expect_sleep_calls([])
        expect(result).to be_nil
      end
    end

    context 'when the timeout is 0' do
      it 'does not yield control' do
        result = nil
        expect do |block|
          result = subject.retry_until_truthy(timeout: -1, &block)
        end.to_not yield_control
        expect_sleep_calls([])
        expect(result).to be_nil
      end
    end

    context 'when the timeout is 40' do
      let(:timeout) { 40 }

      it 'only yields once if the block takes longer than the allocated time' do
        expect do |block|
          subject.retry_until_truthy(timeout: timeout) do
            block.to_proc.call
            increment_time_by(timeout + 1)
          end
        end.to yield_control.once
        expect_sleep_calls([])
      end

      it 'yields exponentially until the timeout has surpassed' do
        result = nil
        expect do |block|
          result = subject.retry_until_truthy(timeout: timeout, &block)
        end.to yield_control.exactly(5).times
        expect_sleep_calls([2, 4, 8, 16, 10])
        expect(result).to be_nil
      end

      it 'returns the yielded value if it is immediately truthy' do
        result = nil
        expect do |block|
          result = subject.retry_until_truthy(timeout: timeout) do
            block.to_proc.call
            :success
          end
        end.to yield_control.exactly(1).times
        expect_sleep_calls([])
        expect(result).to eq(:success)
      end

      it 'returns the yielded value if it is eventually truthy' do
        result = nil
        expect do |block|
          result = subject.retry_until_truthy(timeout: timeout) do
            block.to_proc.call
            mock_clock.now >= 8
          end
        end.to yield_control.exactly(4).times
        expect_sleep_calls([2, 4, 8])
        expect(result).to eq(true)
      end

      it 'raises any unhandled exceptions' do
        error = StandardError.new
        expect do
          subject.retry_until_truthy(timeout: timeout) do
            raise error
          end
        end.to raise_error(error)
        expect(subject).to_not have_received(:sleep)
      end
    end
  end
end
